from pylab import *
from camori_aux import step_renderer, step_get_colors

import scipy.optimize

if __name__ == '__main__':

    ion()


    coords = array([[-.4,1.8],
                    [2.0,-2.9],
                    [-3,4],
                    [1.4,1.7]])
    stepL = []


    figure(1, figsize=(6,6))
    for n, p in enumerate(coords):
        c1, c2 =  np.random.randint(0,256, size=(2,3))
        step = zeros((5,5,3), dtype=np.int)
        step_renderer(p[0], p[1], c1, c2, step)
        stepL.append(step)



    frma = np.load('pic1step.npz')['arr_0']
    frm = zeros((5,5,3), dtype=np.int)
    for k in range(5):
        frm[:,k,:] = frma[k,:,:]

    frmhat = zeros((5,5,3), dtype=np.int)

    # p = [1.0, -.5]
    c1 = copy(frm[2,0])
    c2 = copy(frm[2,4])
    # p = [2.5,-2.5]
    p = [0.0,0.0]
    step_renderer(p[0], p[1], c1, c2, frmhat)

    frmwrk = zeros((5,5,3), dtype=np.int)
    def lossr(x, c1,c2,frm):
        p = x[:2]
        # c1 = array(x[2:5], dtype=np.int)
        # c2 = array(x[5:], dtype=np.int)
        step_renderer(p[0], p[1], c1, c2, frmwrk)
        return np.sum(np.abs(frm - frmwrk))

    p_ini = array(p)

    print p_ini
    print c1, c2

    p_opt = scipy.optimize.fmin_powell(lossr, p_ini, args=(c1,c2,frm,))
    step_get_colors(p_opt[0], p_opt[1], c1, c2, frmwrk)
    print p_opt
    print c1, c2

    p_opt = scipy.optimize.fmin(lossr, p_opt, args=(c1,c2,frm,))
    step_get_colors(p_opt[0], p_opt[1], c1, c2, frmwrk)
    print p_opt
    print c1,c2

    p = p_opt
    
    step_renderer(p[0], p[1], c1, c2, frmwrk)




    figure(1, figsize=(6,6))
    for n, step in enumerate(stepL):
        p = coords[n]
        subplot(2,2,n+1)
        imshow(step/256., extent=(-2.5,2.5,2.5,-2.5))
        plot([p[0], p[1]],[-2.5,2.5] , 'r-')
        #axis([-0.5,5.5,5.5,-0.5])
        axis([-2.5,2.5,2.5,-2.5])

    figure(3, figsize=(6,6))
    subplot(2,2,1)
    imshow(frm/256., extent=(-2.5,2.5,2.5,-2.5))
    subplot(2,2,2)
    imshow(frmhat/256., extent=(-2.5,2.5,2.5,-2.5))
    plot([p_ini[0], p_ini[1]],[-2.5,2.5] , 'r-')

    subplot(2,2,3)
    imshow(frm/256., extent=(-2.5,2.5,2.5,-2.5))
    plot([p_opt[0], p_opt[1]],[-2.5,2.5] , 'r-')
    subplot(2,2,4)
    imshow(frmwrk/256., extent=(-2.5,2.5,2.5,-2.5))
    plot([p_opt[0], p_opt[1]],[-2.5,2.5] , 'r-')
